import requests
import re
import logging
from datetime import datetime, timedelta
from bs4 import BeautifulSoup
from typing import List, Dict, Optional, Tuple
import pymysql
from sqlalchemy import create_engine, text
from sqlalchemy.orm import sessionmaker

logger = logging.getLogger(__name__)

class LotteryCrawler:
    """双色球数据爬取服务"""
    
    def __init__(self):
        self.base_url = "http://kaijiang.zhcw.com/zhcw/html/ssq/list_1.html"
        self.headers = {
            'User-Agent': 'Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.124 Safari/537.36',
            'Accept': 'text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,*/*;q=0.8',
            'Accept-Language': 'zh-CN,zh;q=0.9,en;q=0.8',
            'Accept-Encoding': 'gzip, deflate',
            'Connection': 'keep-alive',
            'Upgrade-Insecure-Requests': '1'
        }
        self.session = requests.Session()
        self.session.headers.update(self.headers)
        
        # 数据库连接配置
        self.db_config = {
            'host': 'localhost',
            'port': 3306,
            'user': 'root',
            'password': 'mario',
            'database': 'ai_ssq',
            'charset': 'utf8mb4'
        }
        
        # 创建数据库引擎
        self.engine = create_engine(
            f"mysql+pymysql://{self.db_config['user']}:{self.db_config['password']}@{self.db_config['host']}:{self.db_config['port']}/{self.db_config['database']}?charset={self.db_config['charset']}"
        )
        self.Session = sessionmaker(bind=self.engine)
    
    def crawl_lottery_data(self, max_pages: int = 3, target_count: int = 30) -> Dict[str, any]:
        """
        爬取双色球数据
        
        Args:
            max_pages: 最大爬取页数，默认3页
            target_count: 目标获取数据条数，默认30条
            
        Returns:
            爬取结果字典
        """
        try:
            logger.info(f"开始爬取双色球数据，目标数量: {target_count}条")
            
            all_data = []
            current_page = 1
            total_added = 0
            total_skipped = 0
            
            while current_page <= max_pages and len(all_data) < target_count:
                logger.info(f"正在爬取第 {current_page} 页...")
                
                # 构建页面URL
                if current_page == 1:
                    page_url = self.base_url
                else:
                    page_url = f"http://kaijiang.zhcw.com/zhcw/html/ssq/list_{current_page}.html"
                
                # 爬取当前页面
                page_data = self._crawl_single_page(page_url)
                
                if not page_data:
                    logger.warning(f"第 {current_page} 页没有获取到数据")
                    break
                
                # 处理页面数据
                for item in page_data:
                    if len(all_data) >= target_count:
                        break
                    
                    # 检查是否已存在
                    existing = self._check_existing_draw(item['draw_number'])
                    
                    if existing:
                        logger.debug(f"期号 {item['draw_number']} 已存在，跳过")
                        total_skipped += 1
                        continue
                    
                    # 验证数据有效性
                    if self._validate_lottery_data(item):
                        all_data.append(item)
                        total_added += 1
                        logger.debug(f"添加期号 {item['draw_number']} 的数据")
                    else:
                        logger.warning(f"期号 {item['draw_number']} 数据无效，跳过")
                        total_skipped += 1
                
                current_page += 1
            
            # 保存到数据库
            saved_count = self._save_to_database(all_data)
            
            logger.info(f"爬取完成: 新增 {saved_count} 条，跳过 {total_skipped} 条")
            
            return {
                'success': True,
                'total_crawled': len(all_data),
                'total_added': saved_count,
                'total_skipped': total_skipped,
                'pages_crawled': current_page - 1,
                'message': f'成功爬取 {saved_count} 条新数据，跳过 {total_skipped} 条重复数据'
            }
            
        except Exception as e:
            logger.error(f"爬取双色球数据失败: {str(e)}")
            return {
                'success': False,
                'error': str(e),
                'total_crawled': 0,
                'total_added': 0,
                'total_skipped': 0,
                'pages_crawled': 0,
                'message': f'爬取失败: {str(e)}'
            }
    
    def _check_existing_draw(self, draw_number: str) -> bool:
        """检查期号是否已存在"""
        try:
            with self.engine.connect() as conn:
                result = conn.execute(
                    text("SELECT COUNT(*) FROM lottery_results WHERE draw_number = :draw_number"),
                    {"draw_number": draw_number}
                )
                count = result.scalar()
                return count > 0
        except Exception as e:
            logger.error(f"检查期号存在性失败: {str(e)}")
            return False
    
    def _crawl_single_page(self, url: str) -> List[Dict]:
        """爬取单个页面的数据"""
        try:
            logger.debug(f"请求URL: {url}")
            
            response = self.session.get(url, timeout=15)
            response.raise_for_status()
            response.encoding = 'utf-8'
            
            soup = BeautifulSoup(response.text, 'html.parser')
            
            # 查找开奖数据表格 - 使用正确的类名
            table = soup.find('table', class_='wqhgt')
            if not table:
                logger.warning("未找到开奖数据表格")
                return []
            
            rows = table.find_all('tr')[2:]  # 跳过前两行表头
            page_data = []
            
            for row in rows:
                try:
                    item = self._parse_table_row(row)
                    if item:
                        page_data.append(item)
                except Exception as e:
                    logger.error(f"解析表格行失败: {str(e)}")
                    continue
            
            logger.debug(f"页面解析完成，获取 {len(page_data)} 条数据")
            return page_data
            
        except requests.exceptions.RequestException as e:
            logger.error(f"请求页面失败: {str(e)}")
            return []
        except Exception as e:
            logger.error(f"解析页面失败: {str(e)}")
            return []
    
    def _parse_table_row(self, row) -> Optional[Dict]:
        """解析表格行数据"""
        try:
            cells = row.find_all('td')
            if len(cells) < 4:
                return None
            
            # 解析开奖日期
            date_text = cells[0].get_text(strip=True)
            if not date_text or date_text == '开奖日期':
                return None
            
            try:
                draw_date = datetime.strptime(date_text, '%Y-%m-%d').date()
            except ValueError:
                logger.warning(f"日期格式错误: {date_text}")
                return None
            
            # 解析期数
            draw_number = cells[1].get_text(strip=True)
            if not draw_number or not draw_number.isdigit():
                return None
            
            # 解析开奖号码
            numbers_text = cells[2].get_text(strip=True)
            
            # 处理连续字符串格式，每两个字符分割一次
            if len(numbers_text) == 14:  # 6个红球 + 1个蓝球 = 14个字符
                numbers = []
                for i in range(0, 14, 2):
                    numbers.append(int(numbers_text[i:i+2]))
            else:
                # 尝试正则表达式提取
                numbers = re.findall(r'\d+', numbers_text)
            
            if len(numbers) != 7:
                logger.warning(f"号码格式错误: {numbers_text}")
                return None
            
            # 验证号码范围
            red_balls = [int(num) for num in numbers[:6]]
            blue_ball = int(numbers[6])
            
            if not all(1 <= ball <= 33 for ball in red_balls):
                logger.warning(f"红球号码超出范围: {red_balls}")
                return None
            
            if not 1 <= blue_ball <= 16:
                logger.warning(f"蓝球号码超出范围: {blue_ball}")
                return None
            
            # 检查红球是否有重复
            if len(set(red_balls)) != 6:
                logger.warning(f"红球号码重复: {red_balls}")
                return None
            
            # 获取奖金信息
            prize_info = cells[3].get_text(strip=True) if len(cells) > 3 else ''
            
            return {
                'draw_date': draw_date,
                'draw_number': draw_number,
                'red_ball_1': red_balls[0],
                'red_ball_2': red_balls[1],
                'red_ball_3': red_balls[2],
                'red_ball_4': red_balls[3],
                'red_ball_5': red_balls[4],
                'red_ball_6': red_balls[5],
                'blue_ball': blue_ball,
                'prize_info': prize_info
            }
            
        except Exception as e:
            logger.error(f"解析表格行异常: {str(e)}")
            return None
    
    def _validate_lottery_data(self, data: Dict) -> bool:
        """验证双色球数据的有效性"""
        try:
            # 检查必要字段
            required_fields = ['draw_date', 'draw_number', 'red_ball_1', 'red_ball_2', 
                             'red_ball_3', 'red_ball_4', 'red_ball_5', 'red_ball_6', 'blue_ball']
            
            for field in required_fields:
                if field not in data or data[field] is None:
                    logger.warning(f"缺少必要字段: {field}")
                    return False
            
            # 验证期数格式
            if not data['draw_number'].isdigit() or len(data['draw_number']) < 7:
                logger.warning(f"期数格式错误: {data['draw_number']}")
                return False
            
            # 验证红球范围
            red_balls = [data['red_ball_1'], data['red_ball_2'], data['red_ball_3'],
                        data['red_ball_4'], data['red_ball_5'], data['red_ball_6']]
            
            if not all(1 <= ball <= 33 for ball in red_balls):
                logger.warning(f"红球号码超出范围: {red_balls}")
                return False
            
            # 验证红球不重复
            if len(set(red_balls)) != 6:
                logger.warning(f"红球号码重复: {red_balls}")
                return False
            
            # 验证蓝球范围
            if not 1 <= data['blue_ball'] <= 16:
                logger.warning(f"蓝球号码超出范围: {data['blue_ball']}")
                return False
            
            # 验证日期合理性
            if data['draw_date'] > datetime.now().date():
                logger.warning(f"开奖日期不能是未来日期: {data['draw_date']}")
                return False
            
            return True
            
        except Exception as e:
            logger.error(f"数据验证异常: {str(e)}")
            return False
    
    def _save_to_database(self, data_list: List[Dict]) -> int:
        """保存数据到数据库"""
        saved_count = 0
        
        try:
            with self.engine.connect() as conn:
                for data in data_list:
                    try:
                        # 再次检查是否已存在（防止并发情况）
                        existing = self._check_existing_draw(data['draw_number'])
                        if existing:
                            logger.debug(f"期号 {data['draw_number']} 已存在，跳过保存")
                            continue
                        
                        # 插入新记录
                        conn.execute(
                            text("""
                                INSERT INTO lottery_results 
                                (draw_date, draw_number, red_ball_1, red_ball_2, red_ball_3, 
                                 red_ball_4, red_ball_5, red_ball_6, blue_ball, prize_info, created_at)
                                VALUES (:draw_date, :draw_number, :red_ball_1, :red_ball_2, :red_ball_3,
                                        :red_ball_4, :red_ball_5, :red_ball_6, :blue_ball, :prize_info, NOW())
                            """),
                            {
                                'draw_date': data['draw_date'],
                                'draw_number': data['draw_number'],
                                'red_ball_1': data['red_ball_1'],
                                'red_ball_2': data['red_ball_2'],
                                'red_ball_3': data['red_ball_3'],
                                'red_ball_4': data['red_ball_4'],
                                'red_ball_5': data['red_ball_5'],
                                'red_ball_6': data['red_ball_6'],
                                'blue_ball': data['blue_ball'],
                                'prize_info': data.get('prize_info', '')
                            }
                        )
                        saved_count += 1
                        
                    except Exception as e:
                        logger.error(f"保存期号 {data['draw_number']} 失败: {str(e)}")
                        continue
                
                # 提交事务
                conn.commit()
                logger.info(f"成功保存 {saved_count} 条数据到数据库")
            
        except Exception as e:
            logger.error(f"数据库保存失败: {str(e)}")
            raise
        
        return saved_count
    
    def get_latest_draw_info(self) -> Optional[Dict]:
        """获取最新一期开奖信息"""
        try:
            with self.engine.connect() as conn:
                result = conn.execute(
                    text("""
                        SELECT draw_number, draw_date, red_ball_1, red_ball_2, red_ball_3,
                               red_ball_4, red_ball_5, red_ball_6, blue_ball, prize_info
                        FROM lottery_results 
                        ORDER BY draw_date DESC 
                        LIMIT 1
                    """)
                )
                row = result.fetchone()
                
                if row:
                    return {
                        'draw_number': row.draw_number,
                        'draw_date': row.draw_date.isoformat(),
                        'red_balls': [row.red_ball_1, row.red_ball_2, row.red_ball_3,
                                     row.red_ball_4, row.red_ball_5, row.red_ball_6],
                        'blue_ball': row.blue_ball,
                        'prize_info': row.prize_info
                    }
                return None
        except Exception as e:
            logger.error(f"获取最新开奖信息失败: {str(e)}")
            return None
    
    def get_database_stats(self) -> Dict:
        """获取数据库统计信息"""
        try:
            with self.engine.connect() as conn:
                # 获取总记录数
                count_result = conn.execute(text("SELECT COUNT(*) FROM lottery_results"))
                total_count = count_result.scalar()
                
                # 获取日期范围
                date_range_result = conn.execute(
                    text("""
                        SELECT MIN(draw_date) as start_date, MAX(draw_date) as end_date
                        FROM lottery_results
                    """)
                )
                date_range = date_range_result.fetchone()
                
                return {
                    'total_records': total_count,
                    'date_range': {
                        'start_date': date_range.start_date.isoformat() if date_range.start_date else None,
                        'end_date': date_range.end_date.isoformat() if date_range.end_date else None
                    },
                    'latest_draw': self.get_latest_draw_info()
                }
        except Exception as e:
            logger.error(f"获取数据库统计失败: {str(e)}")
            return {
                'total_records': 0,
                'date_range': {'start_date': None, 'end_date': None},
                'latest_draw': None
            }
